use std::path::{Path, PathBuf};

use embedding_model::EmbeddingModel;

use crate::{
    content::{ContentCollection, ContentStore},
    SearchConfig, SearchResultItem,
};

pub trait EngineBuilder {
    type InitConfig: Default + Sized;
    type Engine: Engine;
    fn init(config: &Self::InitConfig) -> anyhow::Result<Self>
    where
        Self: Sized;
    fn build(self) -> anyhow::Result<Self::Engine>;
}

pub trait Engine {
    async fn search(
        &self,
        inputs: &str,
        contents: &ContentStore,
        config: &SearchConfig,
    ) -> anyhow::Result<ContentCollection> {
        let ids = self.search_for_ids(inputs, contents, config).await?;
        let result = contents.get_collection(&ids);
        match result {
            Some(s) => Ok(s),
            None => {
                if config.allow_null_result {
                    Ok(ContentCollection::empty())
                } else {
                    Err(anyhow::Error::msg("Null result"))
                }
            }
        }
    }
    async fn search_for_ids(
        &self,
        inputs: &str,
        contents: &ContentStore,
        config: &SearchConfig,
    ) -> anyhow::Result<Vec<SearchResultItem>>;
}

pub use embedding_model::{DType, Device, Result};
pub struct VecEngineImpl<V, E> {
    engine: V,
    model: E,
}

pub trait VecEngineBuilder {
    type InitConfig: Default + Sized;
    type Engine: VecEngine;
    fn init(config: &Self::InitConfig) -> anyhow::Result<Self>
    where
        Self: Sized;
    fn build(self) -> anyhow::Result<Self::Engine>;
}
pub trait VecEngine {
    async fn search(
        &self,
        inputs: &[f32],
        contents: &ContentStore,
        config: &SearchConfig,
    ) -> anyhow::Result<ContentCollection> {
        let ids = self.search_for_ids(inputs, contents, config).await?;
        let result = contents.get_collection(&ids);
        match result {
            Some(s) => Ok(s),
            None => {
                if config.allow_null_result {
                    Ok(ContentCollection::empty())
                } else {
                    Err(anyhow::Error::msg("Null result"))
                }
            }
        }
    }
    async fn search_for_ids(
        &self,
        inputs: &[f32],
        contents: &ContentStore,
        config: &SearchConfig,
    ) -> anyhow::Result<Vec<SearchResultItem>>;
}

impl<V, E> Engine for VecEngineImpl<V, E>
where
    V: VecEngine,
    E: EmbeddingModel,
{
    async fn search(
        &self,
        inputs: &str,
        contents: &ContentStore,
        config: &SearchConfig,
    ) -> anyhow::Result<ContentCollection> {
        let inputs = self.model.embedding_single_vec(inputs)?;
        self.engine.search(&inputs, contents, config).await
    }

    async fn search_for_ids(
        &self,
        inputs: &str,
        contents: &ContentStore,
        config: &SearchConfig,
    ) -> anyhow::Result<Vec<SearchResultItem>> {
        let inputs = self.model.embedding_single_vec(inputs)?;
        self.engine.search_for_ids(&inputs, contents, config).await
    }
}

impl<V, E> VecEngineImpl<V, E>
where
    V: VecEngine,
    E: EmbeddingModel,
{
    pub fn new<B, C, P>(
        vecengine_init_config: &C,
        embeddingmodel_path: P,
        embeddingmodel_dtype: embedding_model::DType,
        embeddingmodel_decice: &embedding_model::Device,
    ) -> anyhow::Result<Self>
    where
        B: VecEngineBuilder<InitConfig = C, Engine = V>,
        C: Default + Sized,
        P: AsRef<Path>,
    {
        Ok(Self {
            engine: B::init(vecengine_init_config)?.build()?,
            model: E::load(
                embeddingmodel_path,
                embeddingmodel_dtype,
                embeddingmodel_decice,
            )?,
        })
    }
    pub fn new_with_embedding_model_config<B, C>(
        vecengine_init_config: &C,
        embedding_model_config: &embedding_model::EmbeddingModelConfig,
    ) -> anyhow::Result<Self>
    where
        B: VecEngineBuilder<InitConfig = C, Engine = V>,
        C: Default + Sized,
    {
        Self::new::<B, C, &PathBuf>(
            vecengine_init_config,
            &embedding_model_config.path,
            embedding_model_config.dtype,
            &embedding_model_config.device,
        )
    }
}
